import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import io.reactivex.Completable;
import org.junit.jupiter.api.Test;

public class _05_Vector {
    public static void main(String[] args) {
        QwenEmbeddingModel embeddingModel = QwenEmbeddingModel
                .builder()
                .apiKey("sk-b957278c2e4f48b39d99115499c60285")
                .build();

        Response<Embedding> embed = embeddingModel.embed("你好，我叫胡治纬");
        System.out.println(embed.content().toString());
        System.out.println(embed.content().vector().length);
    }

    @Test
    public void test02() {
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        QwenEmbeddingModel embeddingModel = QwenEmbeddingModel
                .builder()
                .apiKey("sk-b957278c2e4f48b39d99115499c60285")
                .build();

        TextSegment segment1 = TextSegment.from("""
                预订航班：
                - 通过我们的网站或APP预订
                - 预订时需要全额付款
                - 确保个人信息（姓名，ID等）的准确性，因为更正可能会产生25元的费用
                """);
        Embedding embedding1 = embeddingModel.embed(segment1).content();
        embeddingStore.add(embedding1, segment1);

        TextSegment segment2 = TextSegment.from("""
                取消预订：
                - 最晚在航班起飞前48小时取消
                - 取消费用：经济舱75元，豪华经济舱50元，商务舱25元
                - 退款将在7个工作日内处理
                """);
        Embedding embedding2 = embeddingModel.embed(segment2).content();
        embeddingStore.add(embedding2, segment2);

        Embedding queryEmbedding = embeddingModel.embed("退票需要多少钱").content();

        EmbeddingSearchRequest build = EmbeddingSearchRequest
                .builder()
                .queryEmbedding(queryEmbedding)
//                .maxResults(1)
                .build();

        EmbeddingSearchResult<TextSegment> segementEmbeddingSearchResult = embeddingStore.search(build);
        segementEmbeddingSearchResult.matches().forEach(embeddingMatch -> {
            System.out.println(embeddingMatch.score());
            System.out.println(embeddingMatch.embedded().text());
        });
    }
}
